#pragma once

#include <cstdint>
#include <cassert>
#include <iostream>

namespace uzu
{
    template<typename T>
    struct Matrix
    {
        uint32_t rows;
        uint32_t cols;
        T* data;

    public:
        Matrix(uint32_t r, uint32_t c)
        {
            rows = r;
            cols = c;
            data = new T[rows * cols];
        }

        ~Matrix()
        {
            if (data != nullptr) delete[] data;
        }

        T& operator()(int y, int x) { return data[y * cols + x]; }
        const T& operator()(int y, int x) const { return data[y * cols + x]; }

    };

    // matrix multiplication with cpu
    template<typename T>
    Matrix<T> operator*(const Matrix<T>& mat1, const Matrix<T>& mat2)
    {
        assert(mat1.cols == mat2.rows);
        Matrix<T> result(mat1.rows, mat2.cols);
        for (uint32_t y = 0; y < mat1.rows; ++y)
        {
            for (uint32_t x = 0; x < mat2.cols; ++x)
            {
                T sum = 0;
                for (uint32_t k = 0; k < mat1.cols; ++k)
                    sum += mat1(y, k) * mat2(k, x);
                result(y, x) = sum;
            }
        }
        return result;
    }
    
    template<typename T>
    std::ostream& operator<<(std::ostream& os, const Matrix<T>& mat)
    {
        for (uint32_t y = 0; y < mat.rows; ++y)
        {
            for (uint32_t x = 0; x < mat.cols; ++x)
                os << mat(y, x) << ",";
            os << "\n";
        }
        return os;
    }

    typedef Matrix<float> MatrixF;

    // allocate device memory for matrix
    void AllocateDeviceMemory(float** buffer, const MatrixF& matrix, bool copy);
}
